import uuid
from flask import Flask, request, jsonify, render_template
import chromadb
from chromadb.config import Settings
from pathlib import Path
import json
from flasgger import Swagger
import yaml
import os
import requests

app = Flask(__name__)

# Swagger配置
swagger_config = {
    "headers": [],
    "specs": [{
        "endpoint": 'apispec',
        "route": '/apispec.json',
        "rule_filter": lambda rule: True,
        "model_filter": lambda tag: True,
    }],
    "static_url_path": "/flasgger_static",
    "swagger_ui": True,
    "specs_route": "/apidocs/"
}

# 获取当前文件所在目录的绝对路径
current_dir = os.path.dirname(os.path.abspath(__file__))
# 构建 openapi.yaml 的绝对路径
openapi_path = os.path.join(current_dir, 'openapi.yaml')

# 加载自定义的 OpenAPI 规范
with open(openapi_path, 'r', encoding='utf-8') as f:
    template = yaml.safe_load(f)

# 配置 Swagger
swagger = Swagger(app, config=swagger_config, template=template)

# 初始化 ChromaDB 客户端（这里使用 DuckDB+Parquet 存储方式，可持久化存储数据）
# client = chromadb.Client(
#     Settings(
#         chroma_db_impl="duckdb+parquet",
#         persist_directory="./chroma_db"  # 数据存储目录
#     )
# )
client = chromadb.HttpClient(
    host="localhost",  # 改成你的 ChromaDB 服务器 IP
    port=8000         # 改成你的 ChromaDB 服务器端口
)

# Ollama embedding 模型配置
OLLAMA_BASE_URL = "http://localhost:11434"  # Ollama 服务地址
MODEL_NAME = "bge-m3:latest"             # 使用 nomic-embed-text 模型，你也可以换成其他支持 embedding 的模型

def get_embedding(text):
    """使用 Ollama API 获取文本的 embedding"""
    try:
        response = requests.post(
            f"{OLLAMA_BASE_URL}/api/embeddings",
            json={
                "model": MODEL_NAME,
                "prompt": text
            }
        )
        response.raise_for_status()
        return response.json()["embedding"]
    except Exception as e:
        print(f"Error getting embedding: {str(e)}")
        raise

#######################################
# 前端页面（模板）路由
#######################################

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/collection/<collection_name>')
def collection_page(collection_name):
    return render_template('collection.html', collection_name=collection_name)

#######################################
# API：Collection 的增删改查
#######################################

# 查询所有 collections
@app.route('/api/collections', methods=['GET'])
def list_collections():
    # 获取分页和搜索参数
    page = int(request.args.get("page", 1))
    page_size = int(request.args.get("page_size", 10))
    search = request.args.get("search", "").lower()
    
    try:
        collection_names = client.list_collections()
        collections = []
        
        # 获取所有collections的信息
        for name in collection_names:
            # 如果有搜索条件，先过滤
            if search and search not in name.lower():
                continue
                
            col = client.get_collection(name)
            result = col.get()
            dimension = None
            if result.get("embeddings") and len(result["embeddings"]) > 0:
                dimension = len(result["embeddings"][0])
            
            collections.append({
                'name': col.name,
                'dimension': dimension,
                'count': col.count()
            })
        
        # 计算总数和分页
        total = len(collections)
        start = (page - 1) * page_size
        end = start + page_size
        
        return jsonify({
            'collections': collections[start:end],
            'total': total,
            'page': page,
            'page_size': page_size
        })
    except Exception as e:
        return jsonify({"error": str(e)}), 500

# 新建 collection
@app.route('/api/collections', methods=['POST'])
def create_collection():
    data = request.get_json()
    name = data.get('name')
    if not name:
        return jsonify({"error": "Collection name is required"}), 400
    try:
        collection = client.create_collection(name)
        return jsonify({"message": "Collection created", "collection": {"name": collection.name}})
    except Exception as e:
        return jsonify({"error": str(e)}), 500

# 修改 collection（例如重命名操作：复制数据到新集合然后删除旧集合）
@app.route('/api/collections/<collection_name>', methods=['PUT'])
def update_collection(collection_name):
    data = request.get_json()
    new_name = data.get('new_name')
    if not new_name:
        return jsonify({"error": "New collection name is required"}), 400
    try:
        # 获取旧集合
        old_collection = client.get_collection(name=collection_name)
        # 新建新集合
        new_collection = client.create_collection(new_name)
        # 复制旧集合中的所有记录
        result = old_collection.get()
        ids = result.get("ids", [])
        embeddings = result.get("embeddings", [])
        metadatas = result.get("metadatas", [])
        documents = result.get("documents", [])
        if ids:
            new_collection.add(ids=ids, embeddings=embeddings, metadatas=metadatas, documents=documents)
        # 删除旧集合
        client.delete_collection(collection_name)
        return jsonify({"message": "Collection renamed successfully", "collection": {"name": new_name}})
    except Exception as e:
        return jsonify({"error": str(e)}), 500

# 删除 collection
@app.route('/api/collections/<collection_name>', methods=['DELETE'])
def delete_collection(collection_name):
    try:
        client.delete_collection(collection_name)
        return jsonify({"message": "Collection deleted successfully"})
    except Exception as e:
        return jsonify({"error": str(e)}), 500

#######################################
# API：记录的增删改查与检索（在具体 collection 下）
#######################################

# 分页查询记录，支持简单的“搜索”（在 document 字段中模糊匹配）
@app.route('/api/collections/<collection_name>/records', methods=['GET'])
def list_records(collection_name):
    # 获取分页参数
    page = int(request.args.get("page", 1))
    page_size = int(request.args.get("page_size", 10))
    
    # 获取查询参数
    search_type = request.args.get("search_type")
    record_id = request.args.get("id")
    query_text = request.args.get("query_text")
    metadata_filter = request.args.get("metadata_filter")
    document_filter = request.args.get("document_filter")
    
    try:
        collection = client.get_collection(name=collection_name)
        
        # 处理元数据过滤条件
        where_clause = None
        if metadata_filter:
            try:
                filter_dict = json.loads(metadata_filter)
                if isinstance(filter_dict, dict):
                    # 如果是多个条件，转换为 $and 格式
                    if len(filter_dict) > 1:
                        where_clause = {
                            "$and": [
                                {key: {"$eq": value}} if isinstance(value, (str, int, float, bool)) else {key: value}
                                for key, value in filter_dict.items()
                            ]
                        }
                    # 单个条件，直接转换
                    else:
                        key, value = next(iter(filter_dict.items()))
                        where_clause = {
                            key: {"$eq": value} if isinstance(value, (str, int, float, bool)) else value
                        }
            except json.JSONDecodeError:
                return jsonify({"error": "Invalid metadata filter format"}), 400

        # ID 精确查询
        if search_type == 'id' and record_id:
            result = collection.get(
                ids=[record_id],
                include=["metadatas", "documents"]
            )
            total = len(result["ids"])
        
        # 语义查询
        elif search_type == 'hybrid' and query_text:
            try:
                # 使用 Ollama 生成查询文本的 embedding
                query_embedding = get_embedding(query_text)
                
                n_results_requested = int(request.args.get("n_results", 1000))
                query_params = {
                    "query_embeddings": [query_embedding],
                    "n_results": n_results_requested,
                    "include": ["metadatas", "documents", "distances"]
                }
                if where_clause:
                    query_params["where"] = where_clause
                if document_filter:
                    query_params["where_document"] = {"$contains": document_filter}
                    
                result = collection.query(**query_params)
                # 处理嵌套数组结果
                result = {
                    "ids": result["ids"][0] if isinstance(result["ids"][0], list) else result["ids"],
                    "documents": result["documents"][0] if isinstance(result["documents"][0], list) else result["documents"],
                    "metadatas": result["metadatas"][0] if result.get("metadatas") and isinstance(result["metadatas"][0], list) else result.get("metadatas", []),
                    "distances": result["distances"][0] if result.get("distances") and isinstance(result["distances"][0], list) else result.get("distances", [])
                }
                # 手动处理分页
                start_idx = (page - 1) * page_size
                end_idx = start_idx + page_size
                result = {
                    "ids": result["ids"][start_idx:end_idx],
                    "documents": result["documents"][start_idx:end_idx],
                    "metadatas": result["metadatas"][start_idx:end_idx] if result.get("metadatas") else [],
                    "distances": result["distances"][start_idx:end_idx] if result.get("distances") else []
                }
                # 计算符合过滤条件的总记录数，并取 n_results_requested 与过滤总数的最小值
                total_count = get_filtered_count(
                    collection,
                    where=where_clause,
                    where_document={"$contains": document_filter} if document_filter else None
                )
                total = min(total_count, n_results_requested)
            except Exception as e:
                return jsonify({"error": str(e)}), 500
        
        # 普通查询（包括元数据和文档内容过滤）
        else:
            # 使用 get 方法的内置分页
            get_params = {
                "limit": page_size,
                "offset": (page - 1) * page_size,
                "include": ["metadatas", "documents"]
            }
            if where_clause:
                get_params["where"] = where_clause
            if document_filter:
                get_params["where_document"] = {"$contains": document_filter}
                
            result = collection.get(**get_params)
            if where_clause or document_filter:
                # 获取符合条件的总记录数
                total = get_filtered_count(
                    collection,
                    where=where_clause,
                    where_document={"$contains": document_filter} if document_filter else None
                )
            else:
                total = collection.count()  # 获取总记录数用于分页

        # 构建记录列表
        records = [
            {
                "id": doc_id,
                "document": doc,
                "metadata": meta if meta else {},
                "distance": dist if 'distances' in result else None
            }
            for doc_id, doc, meta, dist in zip(
                result["ids"],
                result["documents"],
                result.get("metadatas", [{}] * len(result["ids"])),
                result.get("distances", [None] * len(result["ids"]))
            )
        ]

        return jsonify({
            "records": records,
            "total": total,
            "page": page,
            "page_size": page_size
        })

    except Exception as e:
        print(f"Error in list_records: {str(e)}")  # 添加错误日志
        return jsonify({"error": str(e)}), 500

# 新建记录
@app.route('/api/collections/<collection_name>/records', methods=['POST'])
def create_record(collection_name):
    data = request.get_json()
    record_id = data.get("id", str(uuid.uuid4()))
    document = data.get("document", "")
    # 如果 metadata 不存在或为空，则不包含在参数中
    metadata = data.get("metadata")
    
    if not document:
        return jsonify({"error": "Document content is required"}), 400
        
    try:
        # 使用 Ollama 生成 embedding
        embedding = get_embedding(document)
        
        collection = client.get_collection(name=collection_name)
        
        # 基础参数
        add_params = {
            "ids": [record_id],
            "embeddings": [embedding],
            "documents": [document]
        }
        
        # 只有在 metadata 存在且不为空时才添加
        if metadata:
            add_params["metadatas"] = [metadata]
            
        collection.add(**add_params)
        return jsonify({
            "message": "Record added successfully", 
            "record_id": record_id
        })
    except Exception as e:
        return jsonify({"error": str(e)}), 500

# 更新记录（模拟：先删除再添加）
@app.route('/api/collections/<collection_name>/records/<record_id>', methods=['PUT'])
def update_record(collection_name, record_id):
    data = request.get_json()
    document = data.get("document", "")
    # 如果 metadata 不存在或为空，则不包含在参数中
    metadata = data.get("metadata")
    
    if not document:
        return jsonify({"error": "Document content is required"}), 400
        
    try:
        # 使用 Ollama 生成 embedding
        embedding = get_embedding(document)
        
        collection = client.get_collection(name=collection_name)
        # 先删除旧记录（如果存在）
        try:
            collection.delete(ids=[record_id])
        except Exception:
            pass
            
        # 基础参数
        add_params = {
            "ids": [record_id],
            "embeddings": [embedding],
            "documents": [document]
        }
        
        # 只有在 metadata 存在且不为空时才添加
        if metadata:
            add_params["metadatas"] = [metadata]
            
        collection.add(**add_params)
        return jsonify({
            "message": "Record updated successfully", 
            "record_id": record_id
        })
    except Exception as e:
        return jsonify({"error": str(e)}), 500

# 删除单个记录
@app.route('/api/collections/<collection_name>/records/<record_id>', methods=['DELETE'])
def delete_record(collection_name, record_id):
    try:
        collection = client.get_collection(name=collection_name)
    except Exception as e:
        return jsonify({"error": f"Collection {collection_name} not found"}), 404
    try:
        collection.delete(ids=[record_id])
        return jsonify({"message": "Record deleted successfully", "record_id": record_id})
    except Exception as e:
        return jsonify({"error": str(e)}), 500

# 批量删除记录，传入 JSON 格式的 { "ids": [id1, id2, ...] }
@app.route('/api/collections/<collection_name>/records', methods=['DELETE'])
def batch_delete_records(collection_name):
    data = request.get_json()
    ids = data.get("ids")
    if not ids or not isinstance(ids, list):
        return jsonify({"error": "List of record ids is required"}), 400
    try:
        collection = client.get_collection(name=collection_name)
    except Exception as e:
        return jsonify({"error": f"Collection {collection_name} not found"}), 404
    try:
        collection.delete(ids=ids)
        return jsonify({"message": "Records deleted successfully", "ids": ids})
    except Exception as e:
        return jsonify({"error": str(e)}), 500
    
def get_filtered_count(collection, where=None, where_document=None):
    """获取符合过滤条件的记录总数"""
    try:
        # 只获取 ids 以提高性能
        result = collection.get(
            where=where,
            where_document=where_document,
            include=[]  # 只包含 ids，不包含其他数据以提高性能
        )
        return len(result["ids"])
    except Exception as e:
        print(f"Error getting filtered count: {str(e)}")
        return 0

#######################################
# 启动 Flask 服务
#######################################

if __name__ == '__main__':
    app.run(debug=True)
